package com.sxr.webflux.util;

import com.sxr.file.FileResponse;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpResponse;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.util.Map;

/**
 * 响应式响应写入工具
 * 将业务响应对象写入ServerHttpResponse
 *
 * @author SXR
 * @since 1.0.0
 */
public class ReactiveResponseWriter {

    private static final int BUFFER_SIZE = 8192;

    /**
     * 写入文件响应
     *
     * @param fileResponse 文件响应对象
     * @param response 响应式HTTP响应
     * @return 异步处理结果
     */
    public Mono<Void> writeFileResponse(FileResponse fileResponse, ServerHttpResponse response) {
        // 设置状态码
        response.setStatusCode(HttpStatus.valueOf(fileResponse.getStatusCode()));

        if (!fileResponse.isSuccess()) {
            // 错误响应
            return writeErrorResponse(fileResponse.getStatusCode(),
                                    fileResponse.getErrorMessage(), response);
        }

        // 设置响应头
        setResponseHeaders(fileResponse, response);

        // 设置内容长度
        if (fileResponse.getContentLength() > 0) {
            response.getHeaders().setContentLength(fileResponse.getContentLength());
        }

        // 写入文件内容
        return writeFileContent(fileResponse, response);
    }

    /**
     * 写入JSON响应
     *
     * @param data 数据对象
     * @param response 响应式HTTP响应
     * @return 异步处理结果
     */
    public Mono<Void> writeJsonResponse(Map<String, String> data, ServerHttpResponse response) {
        response.setStatusCode(HttpStatus.OK);
        response.getHeaders().setContentType(MediaType.APPLICATION_JSON);

        // 构建JSON字符串
        String json = buildJsonString(data);

        DataBufferFactory bufferFactory = response.bufferFactory();
        DataBuffer buffer = bufferFactory.wrap(json.getBytes(StandardCharsets.UTF_8));

        return response.writeWith(Mono.just(buffer));
    }

    /**
     * 写入错误响应
     *
     * @param statusCode 状态码
     * @param message 错误消息
     * @param response 响应式HTTP响应
     * @return 异步处理结果
     */
    public Mono<Void> writeErrorResponse(int statusCode, String message, ServerHttpResponse response) {
        response.setStatusCode(HttpStatus.valueOf(statusCode));
        response.getHeaders().setContentType(MediaType.APPLICATION_JSON);

        String errorJson = "{\"success\":false,\"error\":\"" + escapeJson(message) + "\"}";

        DataBufferFactory bufferFactory = response.bufferFactory();
        DataBuffer buffer = bufferFactory.wrap(errorJson.getBytes(StandardCharsets.UTF_8));

        return response.writeWith(Mono.just(buffer));
    }

    /**
     * 设置响应头
     */
    private void setResponseHeaders(FileResponse fileResponse, ServerHttpResponse response) {
        // 设置基本响应头
        fileResponse.getHeaders().forEach((name, value) ->
            response.getHeaders().add(name, value));

        // 设置安全相关的响应头
        response.getHeaders().add("X-Content-Type-Options", "nosniff");
        response.getHeaders().add("X-Frame-Options", "SAMEORIGIN");
        response.getHeaders().add("X-XSS-Protection", "1; mode=block");

        // 设置缓存控制
        if (!response.getHeaders().containsKey("Cache-Control")) {
            response.getHeaders().add("Cache-Control", "private, max-age=3600");
        }
    }

    /**
     * 写入文件内容
     */
    private Mono<Void> writeFileContent(FileResponse fileResponse, ServerHttpResponse response) {
        InputStream inputStream = fileResponse.getInputStream();
        if (inputStream == null) {
            return Mono.empty();
        }

        DataBufferFactory bufferFactory = response.bufferFactory();

        // 创建响应式流来读取文件内容
        Flux<DataBuffer> dataBufferFlux = Flux.<DataBuffer, InputStream>generate(
            () -> inputStream,
            (stream, sink) -> {
                try {
                    byte[] buffer = new byte[BUFFER_SIZE];
                    int bytesRead = stream.read(buffer);

                    if (bytesRead > 0) {
                        DataBuffer dataBuffer = bufferFactory.allocateBuffer(bytesRead);
                        dataBuffer.write(buffer, 0, bytesRead);
                        sink.next(dataBuffer);
                    } else {
                        sink.complete();
                    }
                } catch (IOException e) {
                    sink.error(e);
                }
                return stream;
            },
            stream -> {
                try {
                    stream.close();
                } catch (IOException e) {
                    // 记录日志但不抛出异常
                }
            }
        ).subscribeOn(Schedulers.boundedElastic()); // 使用有界弹性调度器处理阻塞IO

        return response.writeWith(dataBufferFlux);
    }

    /**
     * 构建JSON字符串
     */
    private String buildJsonString(Map<String, String> data) {
        StringBuilder json = new StringBuilder();
        json.append("{");

        boolean first = true;
        for (Map.Entry<String, String> entry : data.entrySet()) {
            if (!first) {
                json.append(",");
            }
            json.append("\"").append(escapeJson(entry.getKey())).append("\":");
            json.append("\"").append(escapeJson(entry.getValue())).append("\"");
            first = false;
        }

        json.append("}");
        return json.toString();
    }

    /**
     * JSON字符串转义
     */
    private String escapeJson(String str) {
        if (str == null) {
            return "";
        }
        return str.replace("\\", "\\\\")
                  .replace("\"", "\\\"")
                  .replace("\n", "\\n")
                  .replace("\r", "\\r")
                  .replace("\t", "\\t")
                  .replace("\b", "\\b")
                  .replace("\f", "\\f");
    }
}
